#include "model.h"
#include "../../lib/SolderingIron_inferencing/src/SolderingIron_inferencing.h"

// 声明全局变量
extern ei_impulse_handle_t& ei_default_impulse;

/**
 * @brief EIModel构造函数
 */
EIModel::EIModel() : initialized(false) {

}

/**
 * @brief 初始化模型
 * @return true 如果初始化成功，否则返回false
 */
bool EIModel::begin() {
    // 初始化EI分类器
    EI_IMPULSE_ERROR result = init_impulse(&ei_default_impulse);    // 初始化打开模型
    if (result != EI_IMPULSE_OK) {
        ei_printf("Failed to initialize impulse (%d)\n", result);
        return false;
    }

    // 打印模型信息
    ei_printf("Initialized Edge Impulse model0\n");
    ei_printf("Model version: %d\n", ei_default_impulse.impulse->deploy_version);
    ei_printf("Model name: %s\n", ei_default_impulse.impulse->project_name);
    ei_printf("Model deploy ID: %d\n", ei_default_impulse.impulse->project_id);
    ei_printf("Input features: %d\n", ei_default_impulse.impulse->nn_input_frame_size);
    ei_printf("Classes: %d\n", ei_default_impulse.impulse->label_count);

    initialized = true;
    return initialized;
}

/**
 * @brief 运行模型分类器
 * @param data 输入数据数组
 * @param size 数据数组大小
 * @param result_label 返回的分类标签索引
 * @param result_score 返回的分类置信度
 * @param debug 是否启用调试模式（默认false）
 * @return 0 如果成功，负值表示错误代码
 */
int EIModel::runClassifier(
    float* data, 
    size_t size, 
    int* result_label,
    float* result_score,
    bool debug
) {
    if (!initialized) {
        if (!begin()) {
            return -1;
        }
    }
    
    // 创建信号结构体
    ei::signal_t signal;
    EI_IMPULSE_ERROR impulse_result = EI_IMPULSE_OK;
    int signal_result;
    
    // 初始化信号结构体
    signal_result = ei::numpy::signal_from_buffer(data, size, &signal);
    if (signal_result != ei::EIDSP_OK) {
        return -2;
    }
    
    // 创建结果结构体
    ei_impulse_result_t result;
    memset(&result, 0, sizeof(ei_impulse_result_t));
    
    // 运行分类器
    impulse_result = (EI_IMPULSE_ERROR)run_classifier(&signal, &result, debug);
    if (impulse_result != EI_IMPULSE_OK) {
        return -3;
    }
    
    // 查找最高置信度的分类结果
    float max_score = 0.0f;
    int max_label = -1;
    
    for (size_t i = 0; i < ei_default_impulse.impulse->label_count; i++) {
        if (result.classification[i].value > max_score) {
            max_score = result.classification[i].value;
            max_label = i;
        }
    }
    
    *result_label = max_label;
    *result_score = max_score;
    
    return 0; // 成功
}

// 获取特征数量
size_t EIModel::getFeatureCount() {
    if (!initialized) {
        return 0;
    }
    return ei_default_impulse.impulse->nn_input_frame_size;
}

// 获取模型版本
const char* EIModel::getVersion() {
    if (!initialized) {
        return nullptr;
    }
    
    static char version_str[16];
    snprintf(version_str, sizeof(version_str), "%d", ei_default_impulse.impulse->deploy_version);
    return version_str;
}

// 获取标签数量
size_t EIModel::getLabelCount() {
    if (!initialized) {
        return 0;
    }
    return ei_default_impulse.impulse->label_count;
}

// 获取标签
const char* EIModel::getLabel(int index) {
    if (!initialized) {
        return nullptr;
    }
    
    if (index < 0 || index >= (int)ei_default_impulse.impulse->label_count) {
        return nullptr;
    }
    
    return ei_default_impulse.impulse->categories[index];
}